#!/usr/bin/env python3

import os
import launch
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument
from launch.substitutions import LaunchConfiguration
from launch_ros.actions import Node
from ament_index_python.packages import get_package_share_directory
from webots_ros2_driver.webots_launcher import WebotsLauncher
from webots_ros2_driver.webots_controller import WebotsController
from webots_ros2_driver.wait_for_controller_connection import WaitForControllerConnection

def generate_launch_description():
    package_dir = get_package_share_directory('my_webots_sim')
    
    # URDF 文件路径
    urdf_file = os.path.join(package_dir, 'urdf', 'my_robot.urdf')
    
    ros2_control_params = os.path.join(package_dir, 'config', 'ros2_control.yml')

    # 启动 Webots (使用ros2_supervisor模式)
    webots = WebotsLauncher(
        world=os.path.join(package_dir, 'worlds', 'my_world.wbt'),
        mode='realtime',
        ros2_supervisor=True
    )

    # Robot State Publisher
    robot_state_publisher = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        output='screen',
        parameters=[{
            'robot_description': open(urdf_file, 'r').read(),
            'publish_frequency': 100.0,
            'use_sim_time': True,
            'frame_prefix': ''
        }]
    )

    # ROS control spawners
    controller_manager_timeout = ['--controller-manager-timeout', '50']
    controller_manager_prefix = 'python.exe' if os.name == 'nt' else ''
    
    # 差速控制器spawner
    diffdrive_controller_spawner = Node(
        package='controller_manager',
        executable='spawner',
        output='screen',
        prefix=controller_manager_prefix,
        arguments=['diffdrive_controller'] + controller_manager_timeout,
        parameters=[{'use_sim_time': True}],
    )
    
    # 关节状态广播器spawner
    joint_state_broadcaster_spawner = Node(
        package='controller_manager',
        executable='spawner',
        output='screen',
        prefix=controller_manager_prefix,
        arguments=['joint_state_broadcaster'] + controller_manager_timeout,
        parameters=[{'use_sim_time': True}],
    )

    # 设置话题重映射
    use_twist_stamped = 'ROS_DISTRO' in os.environ and (os.environ['ROS_DISTRO'] in ['rolling', 'jazzy'])
    if use_twist_stamped:
        remappings = [('/diffdrive_controller/cmd_vel', '/cmd_vel'), ('/diffdrive_controller/odom', '/odom')]
    else:
        remappings = [('/diffdrive_controller/cmd_vel_unstamped', '/cmd_vel'), ('/diffdrive_controller/odom', '/odom')]

    # base_footprint 到 base_link 的坐标变换
    base_footprint_to_base_link = Node(
        package='tf2_ros',
        executable='static_transform_publisher',
        arguments=['0', '0', '0.0587', '0', '0', '0', 'base_footprint', 'base_link']
    )

    # Webots ROS 2 驱动
    robot_driver = WebotsController(
        robot_name='tank',
        parameters=[
            {'robot_description': urdf_file},
            {'set_robot_state_publisher': False},
            {'use_sim_time': True},
            ros2_control_params
        ],
        output='screen',
        remappings=remappings,
        respawn=True
    )

    # 等待控制器连接
    waiting_nodes = WaitForControllerConnection(
        target_driver=robot_driver,
        nodes_to_start=[diffdrive_controller_spawner, joint_state_broadcaster_spawner]
    )
    
    return LaunchDescription([
        webots,
        webots._supervisor,
        robot_state_publisher,
        base_footprint_to_base_link,
        robot_driver,
        waiting_nodes,
        
        # 当Webots退出时关闭所有节点
        launch.actions.RegisterEventHandler(
            event_handler=launch.event_handlers.OnProcessExit(
                target_action=webots,
                on_exit=[launch.actions.EmitEvent(event=launch.events.Shutdown())]
            )
        )
    ])